from dev.stock_information.stock_detail import get_daily_stock,get_aStock_zh_a_minute,get_week_stock,get_month_stock
from dev.stock_information.stock_detail import get_stock_MACD,get_stock_RSI,get_Stock_Bollinger,get_stock_OBV
from dev.stock_visual.stock_mpf_plot import MACD_plot,RSI_plot,Bollinger_plot,OBV_plot,KLine_plot
import pandas as pd
import os
import datetime

#股票K线指标数据
class Kline_Indices:
    def __init__(self,kline_data,title="",savePath="",indices_params={},indices_flag_dict={}):
        '''
        @param kline_data k线数据（分时/日/周/月K线） type=dataframe，index为DatetimeIndex
        @param indices_params 指标参数字典，比如{"RSI":{"t1" : 6, "t2" : 12, "t3" : 24}, "MACD":{"fastperiod"=12,"slowperiod"=26,"signalperiod"=9}}; type=dict{str:dict{}}
        @param title 图表标题； type=str
        @param savePath 图表保存路径
        @param indices_flag_dict 指标是否生成标志，True则生成指标，False则不生成指标, 比如{ "MACD": True, "RSI": True, "Bollinger": True, "obv": True}； type=dict{str:boolean}
        '''
        self.kline_data = kline_data  #k线数据（分时/日/周/月K线）
        self.title = title  #图表标题
        self.savePath = savePath  #图表保存路径
        self.indices_flag_dict = indices_flag_dict #指标是否生成标志

        self.setting_indices_params = indices_params  #人工设置的参数值

        #各指标的默认参数值
        self.indices_default_params = {
            'MACD' : {'fastperiod':12,'slowperiod':26,'signalperiod':9},
            'RSI' : {'t1':6,'t2':12,'t3':24},
            'Bollinger' : {"timeperiod" : 2, "nbdevup":2, "nbdevdn":2, "matype":0},
            "obv" : {}
        }

        for key in self.indices_default_params:  #使用反射可读性不强，还是使用if-else
            # 判断setting_indices_params指标字段是否为空
            setting_params = self.setting_indices_params[key] if key in self.setting_indices_params.keys() else {}
            default_params = self.indices_default_params[key]
            if(key == 'MACD'):
                fastperiod = setting_params["fastperiod"] if "fastperiod" in setting_params.keys() else default_params["fastperiod"]
                slowperiod = setting_params["slowperiod"] if "slowperiod" in setting_params.keys() else default_params["slowperiod"]
                signalperiod = setting_params["signalperiod"] if "signalperiod" in setting_params.keys() else default_params["signalperiod"]
                self.set_kline_MACD(fastperiod=fastperiod,slowperiod=slowperiod,signalperiod=signalperiod)

            elif(key == 'RSI'):
                t1 = setting_params["t1"] if "t1" in setting_params.keys() else default_params["t1"]
                t2 = setting_params["t2"] if "t2" in setting_params.keys() else default_params["t2"]
                t3 = setting_params["t3"] if "t3" in setting_params.keys() else default_params["t3"]
                self.set_kline_RSI(t1=t1, t2=t2,t3=t3)

            elif(key == 'Bollinger'):
                timeperiod = setting_params["timeperiod"] if "timeperiod" in setting_params.keys() else default_params["timeperiod"]
                nbdevup = setting_params["nbdevup"] if "nbdevup" in setting_params.keys() else default_params["nbdevup"]
                nbdevdn = setting_params["nbdevdn"] if "nbdevdn" in setting_params.keys() else default_params["nbdevdn"]
                matype = setting_params["matype"] if "matype" in setting_params.keys() else default_params["matype"]
                self.set_kline_Bollinger(timeperiod=timeperiod, nbdevup=nbdevup, nbdevdn=nbdevdn, matype=matype)

            elif(key == 'obv'):
                self.set_kline_OBV()

    #设置 MACD指标,包含 macd 白线，macdsignal 黄线，macdhist 柱状图
    def set_kline_MACD(self,fastperiod=12,slowperiod=26,signalperiod=9):
        macd,macdsignal,macdhist = get_stock_MACD(stock_data=self.kline_data['close'],
                                   fastperiod=fastperiod,
                                   slowperiod=slowperiod,
                                   signalperiod=signalperiod)
        self.MACD = (macd,macdsignal,macdhist)
        print(f"MACD = {self.MACD}")

    # 设置RSI指标，包含 RSI6，RSI12, RSI24
    def set_kline_RSI(self,t1=6,t2=12,t3=24):
        RSI6 = get_stock_RSI(stock_data=self.kline_data['close'],timeperiod=t1)
        RSI12 = get_stock_RSI(stock_data=self.kline_data['close'],timeperiod=t2)
        RSI24 = get_stock_RSI(stock_data=self.kline_data['close'],timeperiod=t3)
        self.RSI = (RSI6,RSI12,RSI24)
        print(f"RSI = {self.RSI}")

    #设置 Bollinger指标
    def set_kline_Bollinger(self,timeperiod=2, nbdevup=2, nbdevdn=2, matype=0):
        H_line, M_line, L_line = get_Stock_Bollinger(stock_data=self.kline_data['close'],timeperiod=timeperiod, nbdevup=nbdevup, nbdevdn=nbdevdn, matype=matype)
        self.Bollinger = (H_line, M_line, L_line)
        print(f"Bollinger = {self.Bollinger}")

    #设置obv指标
    def set_kline_OBV(self):
        self.obv = get_stock_OBV(stock_data=self.kline_data)
        print(f"obv = {self.obv}")

    def kline_plot(self):
        '''
        @param  savePath 该股票K线图绘制后保存的图片路径
        @return res 是否绘制成功的返回结果
        '''
        add_plot = []
        pos = 2
        if(self.MACD != None and self.indices_flag_dict["MACD"] == True):
            subplot = MACD_plot(macd=self.MACD[0],macdSignal=self.MACD[1],histogram=self.MACD[2],pos=pos)
            add_plot.extend(subplot)
            pos += 1
        if (self.RSI != None and self.indices_flag_dict["RSI"] == True):
            subplot = RSI_plot(RSI6=self.RSI[0], RSI12=self.RSI[1], RSI24=self.RSI[2], pos=pos)
            add_plot.extend(subplot)
            pos += 1
        if (self.Bollinger != None and self.indices_flag_dict["Bollinger"] == True):
            subplot = Bollinger_plot(H_line=self.Bollinger[0],M_line=self.Bollinger[1], L_line=self.Bollinger[2], pos=pos)
            add_plot.extend(subplot)
            pos += 1
        if (len(self.obv) != 0 and self.indices_flag_dict["obv"] == True):
            subplot = OBV_plot(obv=self.obv, pos=pos)
            add_plot.extend(subplot)
            pos += 1

        res = KLine_plot(self.kline_data,title=self.title,savePath=self.savePath,add_plot=add_plot)
        return res

    #获取图表保存的文件位置
    def get_savePath(self):
        return self.savePath

#单只股票对象
class Stock:
    '''
    每只股票包含：分时K线/日K线/周K线/月K线 4个对象
    '''
    def __init__(self,code,codeName, saveDir):
        self.code = code  #股票代码
        self.codeName = codeName  #股票名称

        self.minute_stock_obj = None  #分时K线数据对象  type=Kline_Indices
        self.daily_stock_obj = None  #日K线数据对象  type=Kline_Indices
        self.week_stock_obj = None   #周K线数据对象  type=Kline_Indices
        self.month_stock_obj = None   #月K线数据对象   type=Kline_Indices

        self.saveDir = saveDir  #该股票信息保存的文件夹位置
        #如果该文件夹不存在会自动创建
        if not os.path.exists(self.saveDir):
            os.mkdir(self.saveDir)

        #K线绘制图片保存的文件路径（分时/日/周/月K）, type=dict
        self.savePaths_dict = {}

        # 标签使用标志
        self.indices_dict = {
            "MACD": True,
            "RSI": True,
            "Bollinger": True,
            "obv": True
        }

    '''设置分时K线对象'''
    def set_minute_K_line(self, period, adjust="",indices_params={},indices_useflag_dict={}):
        '''
        @param: period 滑动窗口大小； type=str
        @param: adjust 默认为空"": 返回不复权的数据，比如qfq返回前复权后的数据，hfq返回后复权后的数据； type=str
        @param: indices_params={} 各项K线指标参数设置，比如{"RSI":{"t1" : 6, "t2" : 12, "t3" : 24}, "MACD":{"fastperiod"=12,"slowperiod"=26,"signalperiod"=9}}; type=dict{str:dict{}}
        @param: indices_useflag_dict={}  要使用那些指标，比如self.indices_dict = { "MACD": True, "RSI": True, "Bollinger": True, "obv": True}
        '''

        # 获取当天的分时数据  type=dataframe
        minute_stock_data = get_aStock_zh_a_minute(symbol=self.code, period=period, adjust=adjust)
        end_date = datetime.datetime.strptime(minute_stock_data.loc[minute_stock_data.index[-1]]["day"],'%Y-%m-%d %H:%M:%S')
        start_date_str = (end_date - datetime.timedelta(days=1)).strftime('%Y-%m-%d %H:%M:%S')
        minute_stock_data = minute_stock_data[minute_stock_data['day'] > start_date_str]

        # 对分时数据中的字符串类型进行强制转换 参考 https://blog.csdn.net/WHYbeHERE/article/details/109497750
        minute_stock_data.timestamp = pd.to_datetime(minute_stock_data.day, format='%Y-%m-%d %H:%M:%S')  # 将day字符串类型强转为datetime
        minute_stock_data.index = minute_stock_data.timestamp  # mpf.kline_plot要求df的index为DatetimeIndex类型
        # 对oepn,high,low,close,volume列进行强制类型转换，参考https://mp.weixin.qq.com/s?__biz=MzU5NzcyMzA5MQ==&mid=2247483928&idx=1&sn=e829a48cfb54720fc44c0f8df6a15d3e&chksm=fe4e50a4c939d9b2260abb027cbb70da022e52ec11a4ec7f1146670b1424a847cbcd3e619edf&scene=27
        for column in minute_stock_data.columns[1:]:
            minute_stock_data[column] = pd.to_numeric(minute_stock_data[column])  # str转numeric

        #将分时K线数据封装成Kline_Indices对象
        title = f"{self.codeName}每{period}分钟的行情数据"
        savePath = os.path.join(self.saveDir, "minute_K_line.png")
        indices_flag_dict = self.indices_dict if len(indices_useflag_dict) == 0 else indices_useflag_dict  #如果indices_useflag_dict为空，则默认全部指标都输出
        self.minute_stock_obj = Kline_Indices(kline_data=minute_stock_data,
                                              indices_params=indices_params,
                                              title=title,
                                              savePath=savePath,
                                              indices_flag_dict=indices_flag_dict)

    '''设置日K线对象'''
    def set_daily_K_line(self,start_date,end_date,adjust="",indices_params={},indices_useflag_dict={}):
        '''
        @param: start_date 起始日期；比如19900101； type=str
        @param: end_date 结束日期；比如21000118； type=str
        @param: adjust 默认为空"": 返回不复权的数据，比如qfq返回前复权后的数据，hfq返回后复权后的数据； type=str
        @param: indices_params={} 各项K线指标参数设置，比如{"RSI":{"t1" : 6, "t2" : 12, "t3" : 24}, "MACD":{"fastperiod"=12,"slowperiod"=26,"signalperiod"=9}}; type=dict{str:dict{}}
        @param: indices_useflag_dict={}  要使用那些指标，比如self.indices_dict = { "MACD": True, "RSI": True, "Bollinger": True, "obv": True}
        '''
        daily_stock_data = get_daily_stock(symbol=self.code, start_date=start_date,end_date=end_date, adjust=adjust)
        daily_stock_data.index = daily_stock_data['date']

        # 将日K线数据封装成Kline_Indices对象
        title = f"{self.codeName}从{start_date}到{end_date}的日K线数据"
        savePath = os.path.join(self.saveDir, "daily_K_line.png")
        indices_flag_dict = self.indices_dict if len(indices_useflag_dict) == 0 else indices_useflag_dict  #如果indices_useflag_dict为空，则默认全部指标都输出
        self.daily_stock_obj = Kline_Indices(kline_data=daily_stock_data,
                                             indices_params=indices_params,
                                             title=title,
                                             savePath=savePath,
                                             indices_flag_dict=indices_flag_dict)

    '''设置周K线对象'''
    def set_week_K_line(self,start_date,end_date,adjust="",indices_params={},indices_useflag_dict={}):
        '''
        @param: start_date 起始日期；比如19900101； type=str
        @param: end_date 结束日期；比如21000118； type=str
        @param: adjust 默认为空"": 返回不复权的数据，比如qfq返回前复权后的数据，hfq返回后复权后的数据； type=str
        @param: indices_params={} 各项K线指标参数设置，比如{"RSI":{"t1" : 6, "t2" : 12, "t3" : 24}, "MACD":{"fastperiod"=12,"slowperiod"=26,"signalperiod"=9}}; type=dict{str:dict{}}
        @param: indices_useflag_dict={}  要使用那些指标，比如self.indices_dict = { "MACD": True, "RSI": True, "Bollinger": True, "obv": True}

        周线就是指一周的K线，以周一开盘价（如果周一不开盘，就从周二开始，以此类推）为周线的开盘价，以周五收盘价为周线收盘价。以一周内的最高价为周线最高价，以一周内的最低价为周线最低价。
        '''
        daily_stock_data = get_daily_stock(symbol=self.code, start_date=start_date, end_date=end_date,
                                                adjust=adjust)
        daily_stock_data.index = daily_stock_data['date']
        week_stock_data = get_week_stock(daily_stock_data)

        # 将周K线数据封装成Kline_Indices对象
        title = f"{self.codeName}从{start_date}到{end_date}的周K线数据"
        savePath = os.path.join(self.saveDir,"week_K_line.png")
        indices_flag_dict = self.indices_dict if len(indices_useflag_dict) == 0 else indices_useflag_dict  #如果indices_useflag_dict为空，则默认全部指标都输出
        self.week_stock_obj = Kline_Indices(kline_data=week_stock_data,
                                            indices_params=indices_params,
                                            title=title,
                                            savePath=savePath,
                                            indices_flag_dict=indices_flag_dict)

    '''设置月K线对象'''
    def set_month_K_line(self,start_date,end_date,adjust="",indices_params={},indices_useflag_dict={}):
        '''
        @param: start_date 起始日期；比如19900101； type=str
        @param: end_date 结束日期；比如21000118； type=str
        @param: adjust 默认为空"": 返回不复权的数据，比如qfq返回前复权后的数据，hfq返回后复权后的数据； type=str
        @param: indices_params={} 各项K线指标参数设置，比如{"RSI":{"t1" : 6, "t2" : 12, "t3" : 24}, "MACD":{"fastperiod"=12,"slowperiod"=26,"signalperiod"=9}}; type=dict{str:dict{}}
        @param: indices_useflag_dict={}  要使用那些指标，比如self.indices_dict = { "MACD": True, "RSI": True, "Bollinger": True, "obv": True}

        在股市中，一般以一个日期为中心，在这个日期经过一个月后，就会算一个周期，长久下来，几个周期逐渐组成了月线。
        '''
        daily_stock_data = get_daily_stock(symbol=self.code, start_date=start_date, end_date=end_date,
                                           adjust=adjust)
        daily_stock_data.index = daily_stock_data['date']
        month_stock_data = get_month_stock(daily_stock_data)

        # 将月K线数据封装成Kline_Indices对象
        title = f"{self.codeName}从{start_date}到{end_date}的月K线数据"
        savePath = os.path.join(self.saveDir, "month_K_line.png")
        indices_flag_dict = self.indices_dict if len(indices_useflag_dict) == 0 else indices_useflag_dict  #如果indices_useflag_dict为空，则默认全部指标都输出
        self.month_stock_obj = Kline_Indices(kline_data=month_stock_data,
                                            indices_params=indices_params,
                                            title=title,
                                            savePath=savePath,
                                            indices_flag_dict=indices_flag_dict)

    #K线绘图：分时K线、日K线、周K线、月K线 (包括MA线）
    def plot_K_line(self):
        #绘制分时K线
        if(self.minute_stock_obj != None):
            res = self.minute_stock_obj.kline_plot()  #res=True则绘制成功，否则失败
            if res == True:
                self.savePaths_dict['minute_kline_path'] = self.minute_stock_obj.get_savePath()
        #绘制日K线
        if (self.daily_stock_obj != None):
            res = self.daily_stock_obj.kline_plot() #res=True则绘制成功，否则失败
            if res == True:
                self.savePaths_dict['daily_kline_path'] = self.daily_stock_obj.get_savePath()
        #绘制周K线
        if (self.week_stock_obj != None):
            res = self.week_stock_obj.kline_plot() #res=True则绘制成功，否则失败
            if res == True:
                self.savePaths_dict['week_kline_path'] = self.week_stock_obj.get_savePath()
        #绘制月K线
        if (self.month_stock_obj != None):
            res = self.month_stock_obj.kline_plot() #res=True则绘制成功，否则失败
            if res == True:
                self.savePaths_dict['month_kline_path'] = self.month_stock_obj.get_savePath()

    #返回该股票的绘制图片路径集合
    def get_savePaths_dict(self):
        return self.savePaths_dict